import math
from dataclasses import dataclass
from typing import Tuple

import cv2
import numpy as np
import open3d as o3d


@dataclass
class CuboidPoint:
    x: float
    y: float
    z: float

    def __init__(self, x: float, y: float, z: float):
        self.x = x
        self.y = y
        self.z = z

    def offset(self, offset_x: float, offset_y: float, offset_z: float):
        self.x += offset_x
        self.y += offset_y
        self.z += offset_z


@dataclass
class CuboidSurface:
    top_left: CuboidPoint
    top_right: CuboidPoint
    bottom_left: CuboidPoint
    bottom_right: CuboidPoint

    def __init__(self, top_left: CuboidPoint, top_right: CuboidPoint,
                 bottom_left: CuboidPoint, bottom_right: CuboidPoint):
        self.top_left = top_left
        self.top_right = top_right
        self.bottom_left = bottom_left
        self.bottom_right = bottom_right

    def offset(self, offset_x: float, offset_y: float, offset_z: float):
        self.top_left.offset(offset_x, offset_y, offset_z)
        self.top_right.offset(offset_x, offset_y, offset_z)
        self.bottom_left.offset(offset_x, offset_y, offset_z)
        self.bottom_right.offset(offset_x, offset_y, offset_z)


def rotation(x: float, y: float, theta: float) -> Tuple[float, float]:
    """rotation in 2d
    """
    theta = theta * math.pi / 180
    a = x * math.cos(theta) + y * math.sin(theta)
    b = y * math.cos(theta) - x * math.sin(theta)
    return a, b


def pass_through_filter(pcd: o3d.geometry.PointCloud, x_range: Tuple[int, int],
                        y_range: Tuple[int, int],
                        z_range: Tuple[int, int]) -> o3d.geometry.PointCloud:
    points = np.asarray(pcd.points)

    x_flag = np.logical_and(points[:, 0] >= x_range[0],
                            points[:, 0] <= x_range[1])
    y_flag = np.logical_and(points[:, 1] >= y_range[0],
                            points[:, 1] <= y_range[1])
    z_flag = np.logical_and(points[:, 2] >= z_range[0],
                            points[:, 2] <= z_range[1])
    flag = np.logical_and(np.logical_and(x_flag, y_flag), z_flag)

    pcd.points = o3d.utility.Vector3dVector(points[flag])

    return pcd


def convert_pcd_to_image(pcd: o3d.geometry.PointCloud,
                         expand_size: int = 100) -> Tuple[np.array, int, int]:
    """将3D点云投影成2D图像

    Parameters
    ----------
    pcd : o3d.geometry.PointCloud
    expand_size : int, optional
        投影成2D图像时, 图像扩展的尺寸, 便于计算轮廓. by default 100

    Returns
    -------
    Tuple[np.array, int, int]
        投影出的2D图像, x方向上的偏移:offset_x, y方向上的偏移: offset_y
        dst_x = src_x + offset_x
        dst_y = src_y + offset_y
    """
    points = np.asarray(pcd.points)

    x_min, x_max = int(np.min(points[:, 0])), int(np.max(points[:, 0]))
    y_min, y_max = int(np.min(points[:, 1])), int(np.max(points[:, 1]))
    img_width = x_max - x_min + expand_size
    img_height = y_max - y_min + expand_size
    img = np.zeros(shape=(img_height, img_width), dtype=np.uint8)

    mask = points[:, :-1] - [x_min, y_min] + \
        [expand_size // 2, expand_size // 2]
    mask = mask.astype(int)
    x_mask, y_mask = mask[:, 0], mask[:, 1]
    img[y_mask, x_mask] = 255

    offset_x = -1 * x_min + expand_size // 2
    offset_y = -1 * y_min + expand_size // 2

    return img, offset_x, offset_y


def get_largest_contour(image: np.ndarray, spot_distance: int) -> np.ndarray:
    kernel = np.ones((spot_distance, spot_distance), dtype=np.uint8)
    dilated_img = cv2.dilate(image, kernel, 1)

    # debug
    # origin_img_show = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    # dilated_img_show = cv2.cvtColor(dilated_img, cv2.COLOR_GRAY2BGR)
    # cv2.imshow("origin_img_show", origin_img_show)
    # cv2.imshow("dilated_img_show", dilated_img_show)
    # cv2.waitKey()
    # debug

    contours, _ = cv2.findContours(dilated_img, cv2.RETR_LIST,
                                   cv2.CHAIN_APPROX_NONE)
    contours = sorted(contours, key=lambda x: cv2.contourArea(x), reverse=True)
    return contours[0]


def cuboid_locate(
        pcd: o3d.geometry.PointCloud, \
        x_range: Tuple[int, int], \
        y_range: Tuple[int, int], \
        z_range: Tuple[int, int], \
        spot_distance: int, \
        cuboid_length: float, \
        cuboid_width: float) -> Tuple[o3d.geometry.PointCloud, CuboidSurface]:
    """集装箱四个角点的定位

    Parameters
    ----------
    pcd : o3d.geometry.PointCloud
        集装箱的3D点云
    x_range : Tuple[int, int]
        x方向上的初始范围
    y_range : Tuple[int, int]
        y方向上的初始范围
    z_range : Tuple[int, int]
        z方向上的初始范围
    spot_distance : int
        相邻光斑之间的距离
    cuboid_length: int
        集装箱的长度
    cuboid_width: int
        集装箱的宽度

    Returns
    -------
    Tuple[o3d.geometry.PointCloud, CuboidSurface]
        经过处理后的点云, 方便调试
        集装箱平面四个角点的位置
    """
    pcd = pass_through_filter(pcd, x_range, y_range, z_range)
    # 1. convert to 2d image
    expand_size = 100
    img, offset_x, offset_y = convert_pcd_to_image(pcd, expand_size)
    # 2. keep only the cuboid points
    contour = get_largest_contour(img, spot_distance)
    img.fill(0)
    cv2.fillPoly(img, [contour], (255, 255, 255))

    # debug start
    # show_img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    # cv2.imshow("first_contour", show_img)
    # cv2.waitKey()
    # debug end

    y_indices, x_indices = np.nonzero(img)
    min_x = np.min(x_indices) - offset_x
    max_x = np.max(x_indices) - offset_x
    min_y = np.min(y_indices) - offset_y
    max_y = np.max(y_indices) - offset_y
    # 3. remove outlier points in cuboid
    pcd = pass_through_filter(pcd, (min_x, max_x), (min_y, max_y), (0, 10000))
    mean_z = np.mean(np.asarray(pcd.points)[:, 2])
    std_z = np.std(np.asarray(pcd.points)[:, 2])
    pcd = pass_through_filter(pcd, (min_x, max_x), (min_y, max_y),
                              (mean_z - std_z * 2, mean_z + std_z * 2))
    # 4. get the cuboid center and rotation
    img, offset_x, offset_y = convert_pcd_to_image(pcd)
    contour = get_largest_contour(img, spot_distance)
    rect = cv2.minAreaRect(contour)
    c_x, c_y = rect[0]
    angle = rect[2]

    # 5. compute the surface corners
    # 此处: x方向对应集装箱的长度, y方向对应集装箱的宽度
    # TODO: 这里的计算方式有待改进, 这里面完全写死了, 实际情况可能需要动态适应, 计算方式需要再考虑下
    half_length, half_width = cuboid_length / 2, cuboid_width / 2

    if angle < 0:
        angle += 90
        top_left_x, top_left_y = rotation(-half_length, -half_width, -angle)
        top_right_x, top_right_y = rotation(half_length, -half_width, -angle)
        bottom_left_x, bottom_left_y = rotation(-half_length, half_width, -angle)
        bottom_right_x, bottom_right_y = rotation(half_length, half_width, -angle)
    else:
        top_left_x, top_left_y = rotation(-half_length, -half_width, angle)
        top_right_x, top_right_y = rotation(half_length, -half_width, -angle)
        bottom_left_x, bottom_left_y = rotation(-half_length, half_width, angle)
        bottom_right_x, bottom_right_y = rotation(half_length, half_width, -angle)

    # debug
    # show_img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    # cv2.circle(show_img, (int(c_x), int(c_y)), 5, (0, 255, 0), -1)
    # cv2.circle(show_img, (int(c_x + top_left_x),     int(c_y + top_left_y)), 5, (0, 255, 0), -1)
    # cv2.circle(show_img, (int(c_x + top_right_x),    int(c_y + top_right_y)), 5, (0, 255, 0), -1)
    # cv2.circle(show_img, (int(c_x + bottom_left_x),  int(c_y + bottom_left_y)), 5, (0, 255, 0), -1)
    # cv2.circle(show_img, (int(c_x + bottom_right_x), int(c_y + bottom_right_y)), 5, (0, 255, 0), -1)
    # cv2.imshow("show", show_img)
    # cv2.waitKey()
    # debug

    # 6. get the cuboid surface
    surface = CuboidSurface(
        CuboidPoint(top_left_x + c_x, top_left_y + c_y, mean_z),
        CuboidPoint(top_right_x + c_x, top_right_y + c_y, mean_z),
        CuboidPoint(bottom_left_x + c_x, bottom_left_y + c_y, mean_z),
        CuboidPoint(bottom_right_x + c_x, bottom_right_y + c_y, mean_z))
    surface.offset(-1 * offset_x, -1 * offset_y, 0)

    # show_img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    # cv2.circle(show_img, (int(c_x), int(c_y)), 5, (0, 255, 0), -1)
    # cv2.circle(show_img, (int(c_x + top_left_x),     int(c_y + top_left_y)), 5, (0, 255, 0), -1)
    # cv2.circle(show_img, (int(c_x + top_right_x),    int(c_y + top_right_y)), 5, (0, 255, 0), -1)
    # cv2.circle(show_img, (int(c_x + bottom_left_x),  int(c_y + bottom_left_y)), 5, (0, 255, 0), -1)
    # cv2.circle(show_img, (int(c_x + bottom_right_x), int(c_y + bottom_right_y)), 5, (0, 255, 0), -1)
    # cv2.imshow("show", show_img)
    # cv2.waitKey()

    return pcd, surface
